#PBAR_GenerateCargoDiagnosticImages.py

# -*- coding: utf-8 -*-
"""
Created on Thu Oct 10 13:59:03 2013

@author: jkwong
"""

#PBAR_LiningUpRadZspecImages.py

import PBAR_Zspec, PBAR_Cargo
import numpy as np
import os, copy, glob
import matplotlib.pyplot as plt
from matplotlib import cm
from scipy import ndimage

# Set useful plot variables
plotColors = ['r', 'b', 'g', 'm', 'c', 'y', 'k'] * 10
lineStyles = ['-', '-.', ':', '_', '|'] *  10
markerTypes = ['.', 'o', 'v', '^', '<', '>', '1', '2', '3', '4', 's', 'p', '*', 'h', 'H', '+', 'x', 'D', 'd']
markerTypes = markerTypes * 2
colorList = np.array(['r', 'b', 'g', 'm', 'c', 'y'])

# list of zpsec detector names - Z1 to Z136
ZspecNameList = np.array(['Z%d'%i for i in np.arange(1, 137)])
# list of bad detector numbers starting at 1
badZspecList = np.array([1,2,3,4,5,6,7,8,20,26,31,33,34,38,39,40,44,53,56,62,68,76,80,125,126,127,128,129,130,131,132,133,134,135,136])
# list of good detector numbers starting at 1
goodZspecList = np.array([i for i in np.arange(1, 137) if (i not in badZspecList)])
# list of bad detector names
badZspecNameList = ZspecNameList[(badZspecList-1)]
# list of good detector names
goodZspecNameList = np.array(['Z%d'%i for i in np.arange(1, 137) if (i not in badZspecList)])

badZspecIndices = badZspecList - 1
goodZspecIndices = goodZspecList - 1

badZspecMask = np.zeros(len(ZspecNameList))
badZspecMask[badZspecIndices] = True
badZspecMask = badZspecMask.astype(bool)

goodZspecMask = np.zeros(len(ZspecNameList))
goodZspecMask[goodZspecIndices] = True
goodZspecMask = goodZspecMask.astype(bool)

dataDir = r'C:\Users\jkwong\Documents\Work\PBAR\data3\BasicScansStandardWidth'

plotDir = r'C:\Users\jkwong\Documents\Work\PBAR\data3\CargoPlotsStandardWidth'

# define parameters for standard width images
numberTimeSlices = 2000
preCargoTimeSlices = 70

acqTime = 1/60.
figureSize = (16, 10)

#################################
## DEFINE FILE NAMES

cargoConfigBaseDirList = glob.glob(os.path.join(dataDir, '*'))

#cargoConfigBaseDirList =cargoConfigBaseDirList[23:]
#cargoConfigBaseDirList = cargoConfigBaseDirList[31:33]
cargoConfigBaseDirList = cargoConfigBaseDirList[34:43]

for (dirIndex, cargoConfigBaseDir) in enumerate(cargoConfigBaseDirList):
    subDirList = glob.glob(os.path.join(cargoConfigBaseDir, '*'))
    for (subDirIndex, dataPath) in enumerate(subDirList):
        print(dataPath)

        a, b = os.path.split(dataPath)
        a, c = os.path.split(a)

        plotSaveDir = os.path.join(plotDir, c, b)
        
        # make paths if not present
        if not os.path.exists(os.path.join(plotDir, c)):
            os.mkdir(os.path.join(plotDir, c))

        if not os.path.exists(plotSaveDir):
            os.mkdir(plotSaveDir)
        
        acquisitionTimeList = []
        zspecScanNumberList = []
        cargoScanNumberList = []
        
        fullFilenameList = []
        filenameList = []
        
        filenameCargoList = []
        fullFilenameCargoList = []
        markerfullFilenameCargoList = []
        
        # find the zspec scan number
        temp = glob.glob(os.path.join(dataPath, '*.npy'))
        # if file doesn't exist skip this set completely
        if temp == []:
            print('No npy file found.  Skipping.')
            continue
        a, filename = os.path.split(temp[0])
        zspecScanNumberList.append(filename[0:4])
        filenameList.append(filename)
        fullFilenameList.append(os.path.join(dataPath, filename))
        
        # get the cargo image name
        temp = glob.glob(os.path.join(dataPath, '*.cargoimage.npy'))
        # if file doesn't exist skip this set completely
        if temp == []:
            print('No cargoimage file found.  Skipping.')
            continue
        a, filename = os.path.split(temp[0])
        cargoScanNumberList.append(filename[0:16])
        filenameCargoList.append(filename)    
        fullFilenameCargoList.append(os.path.join(dataPath, filename))
        
        markerfullFilenameCargoList.append(os.path.join(dataPath, filename.replace('cargoimage.npy', 'cargomarker')))
        acquisitionTimeList.append(acqTime)
        
        
        ######################
        ## READ IN DATA
        
        # Read in marker files
        markerList = []
        for (index, f) in enumerate(markerfullFilenameCargoList):
            markerList.append(PBAR_Cargo.ReadCargoMarker(f))
        
        # Read Cargo
        datCargo = []
        
        for (index, f) in enumerate(fullFilenameCargoList):
#            (A,bpp,formatt,flag,low1,high1,low2,high2) = PBAR_Cargo.ReadCargoImage(f)
            A = np.load(f)            
            datCargo.append(A)


        ###################
        ## PLOTS 
        index = 0
        
        ## Unadjusted Images ##        
        logIntensity = True
        plotMarkers = True
        
        plt.figure(figsize = figureSize)
        plt.grid()
        
        if logIntensity:
            intensity = np.log(datCargo[index].T)
        else:
            intensity = datCargo[index].T
        
        plt.imshow(intensity, interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)
        plt.colorbar()
        plt.title(filenameCargoList[index] + ' Original Size')
        plt.xlim((0, 2400))
        
        if plotMarkers:
            marker = markerList[index]
            for i, mark in enumerate( marker):
                x = mark['rec_left'], mark['rec_right'], mark['rec_right'], mark['rec_left'], mark['rec_left']
                y = mark['rec_bottom'], mark['rec_bottom'], mark['rec_top'], mark['rec_top'], mark['rec_bottom']
                plt.plot(x, y, 'r')
                plt.plot(mark['x'], mark['y'], 'xr')
                if 'left' in mark:
                    plt.plot(mark['left']['x'], mark['left']['y'], 'gx')
                if 'right' in mark:
                    plt.plot(mark['right']['x'], mark['right']['y'], 'go')
                plt.text(max(x), min(y), mark['target'], color = 'g', fontsize = 14)
            
         # Save image to file
        plotFullFilename = os.path.join(plotSaveDir, '%s_OriginalSize.png' %filenameCargoList[index])
        plt.savefig(plotFullFilename)
        plt.close('all')
                    
        plt.figure(figsize = figureSize)
        plt.grid()
        
        
        ###################
        ## PLOTS 
        index = 0
        
        ## Unadjusted Images ##
        logIntensity = False
        plotMarkers = True
        
        plt.figure(figsize = figureSize)
        plt.grid()
        
        if logIntensity:
            intensity = np.log(datCargo[index].T)
        else:
            intensity = datCargo[index].T
        
        plt.imshow(intensity, interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)
        plt.colorbar()
        plt.title(filenameCargoList[index] + ' Original Size')
        plt.xlim((0, 2400))
        
        if plotMarkers:
            marker = markerList[index]
            for i, mark in enumerate( marker):
                x = mark['rec_left'], mark['rec_right'], mark['rec_right'], mark['rec_left'], mark['rec_left']
                y = mark['rec_bottom'], mark['rec_bottom'], mark['rec_top'], mark['rec_top'], mark['rec_bottom']
                plt.plot(x, y, 'r')
                plt.plot(mark['x'], mark['y'], 'xr')
                if 'left' in mark:
                    plt.plot(mark['left']['x'], mark['left']['y'], 'gx')
                if 'right' in mark:
                    plt.plot(mark['right']['x'], mark['right']['y'], 'go')
                plt.text(max(x), min(y), mark['target'], color = 'g', fontsize = 14)
            
         # Save image to file
        plotFullFilename = os.path.join(plotSaveDir, '%s_OriginalSizeLOG.png' %filenameCargoList[index])
        plt.savefig(plotFullFilename)
        plt.close('all')
                    
        plt.figure(figsize = figureSize)
        plt.grid()
        